package com.github.kezhenxu94.data.sink;

import org.apache.flume.Event;
import org.apache.flume.EventDeliveryException;
import org.apache.flume.api.RpcClient;
import org.apache.flume.api.RpcClientFactory;
import org.apache.flume.event.EventBuilder;

import java.util.Map;
import java.util.Properties;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * Created by kezhenxu on 1/7/17.
 *
 * @author kezhenxu (kezhenxu94@163.com)
 */
public class FlumeSink<T> implements Sink<T> {

    protected Properties properties;
    protected RpcClient client;
    protected BlockingQueue<T> queue;
    protected ExecutorService executor;
    protected volatile boolean isRunning;
    protected Function<T, byte[]> valueSupplier;
    protected Function<T, Map<String, String>> headerSupplier;

    @Override
    public void setup() throws Exception {
        properties = new Properties();
        properties.load(getClass().getResourceAsStream("/flume-sink.properties"));
        client = RpcClientFactory.getInstance(properties);
        queue = new LinkedBlockingQueue<>();
        executor = Executors.newSingleThreadExecutor();
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            while (isRunning()) {
                try {
                    T take = queue.take();
                    Event event = EventBuilder.withBody(valueSupplier.apply(take), headerSupplier == null ? null : headerSupplier.apply(take));
                    client.append(event);
                } catch (InterruptedException | EventDeliveryException e) {
                    getLogger().error(e.getMessage(), e);
                }
            }
        });
    }

    @Override
    public void close() throws Exception {
        try {
            isRunning = false;
            while (isRunning()) {
                TimeUnit.MILLISECONDS.sleep(500L);
            }
            executor.awaitTermination(5L, TimeUnit.SECONDS);
        } finally {
            if (client != null) {
                client.close();
            }
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public void put(T t) throws Exception {
        queue.offer(t);
    }

    public Function<T, byte[]> getValueSupplier() {
        return valueSupplier;
    }

    public void setValueSupplier(Function<T, byte[]> valueSupplier) {
        this.valueSupplier = valueSupplier;
    }

    public Function<T, Map<String, String>> getHeaderSupplier() {
        return headerSupplier;
    }

    public void setHeaderSupplier(Function<T, Map<String, String>> headerSupplier) {
        this.headerSupplier = headerSupplier;
    }
}
